from fastapi import FastAPI

from forward.config import get_settings

import asyncio

from fastapi import FastAPI
from fastapi.responses import JSONResponse
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.types import ASGIApp

from forward.http_base import HttpResp



class TimeoutMiddleware(BaseHTTPMiddleware):
    def __init__(self, app: ASGIApp, timeout: int = 15):
        super().__init__(app)
        self.timeout = timeout

    async def dispatch(self, request, call_next):
        try:
            return await asyncio.wait_for(call_next(request), timeout=self.timeout)
        except asyncio.TimeoutError:
            return JSONResponse(
                {
                    "code": HttpResp.SYSTEM_TIMEOUT_ERROR.code,
                    "msg": HttpResp.SYSTEM_TIMEOUT_ERROR.msg,
                    "data": [],
                }
            )


def init_timeout_middleware(app: FastAPI):
    """初始化 超时处理 中间件"""
    app.add_middleware(TimeoutMiddleware, timeout=get_settings().request_timeout)
